import numpy as np
import networkx as nx
import random



def opinion_update(G, node):
  # Fetch neighbors, identity, and opinions of the node
  neighbors = list(G[node])
  identity = G.nodes[node]['identity']
  opinions = G.nodes[node]['opinions']
  beta = G.nodes[node]['beta']

  # Initialize variables for co-partisan and opp-partisan opinions
  opp_opinions = np.zeros_like(opinions)
  opp_count = 0
  co_opinions = np.zeros_like(opinions)
  co_count = 0

  # Aggregate opinions from neighbors based on their identities
  for neighbor in neighbors:
    neighbor_identity = G.nodes[neighbor]['identity']
    neighbor_opinion = G.nodes[neighbor]['opinions']
    if neighbor_identity == identity:
      co_count += 1
      co_opinions += neighbor_opinion
    else: 
      opp_count += 1
      opp_opinions += neighbor_opinion

  # Compute average opinions, handling cases where counts are zero
  if co_count > 0:
    co_opinions /= co_count
  if opp_count > 0:
    opp_opinions /= opp_count
  
  # Update the node's opinions based on beta and neighbors' opinions
  new_opinions = opinions + beta * co_opinions - beta * opp_opinions
  new_opinions = np.tanh(new_opinions)
  G.nodes[node]['opinions'] = new_opinions

def dissim_deg(G, node, target):
    beta = G.nodes[node]['beta']
    node_identity = G.nodes[node]['identity']
    target_identity = G.nodes[target]['identity']
    id_diff = 1 if node_identity != target_identity else 0

    node_opinions = G.nodes[node]['opinions']
    target_opinions = G.nodes[target]['opinions']
    opinion_diff = np.abs(target_opinions - node_opinions)
    opinion_diff_avg = np.mean(opinion_diff)

    dis_degree = beta * id_diff + (1 - beta) * opinion_diff_avg
    return dis_degree

def delete_tie(G, node):
  # Fetch neighbors
  neighbors = list(G[node])
  if not neighbors:
        return  # No neighbors to delete

  # Randomly choose one of the node's neighbors and calculate dissimilarity degree
  target = random.choice(neighbors)
  dis = dissim_deg(G, node, target)
  alpha = G.nodes[node]['alpha']

  # Delete the tie if dissimilarity degree > alpha
  if dis > alpha: 
    G.remove_edge(node, target)
  
  # Delete with 50% chance if dissimilarity degree == alpha
  if dis == alpha:
    if random.random() < 0.5:
      G.remove_edge(node, target)


def add_tie(G, node):
  # Randomly choose a non-neighbor and calculate dissimilarity degree
  non_neighbors = list(nx.non_neighbors(G, node))
  if not non_neighbors:
        return  # No potential nodes to connect
  target = random.choice(non_neighbors)
  dis = dissim_deg(G, node, target)
  alpha = G.nodes[node]['alpha']

  # Add the tie if dissimilarity degree < alpha
  if dis < alpha:
    G.add_edge(node, target)

  # Add with 50% cance if dissimilarity degree == alpha
  if dis == alpha:
    if random.random() < 0.5:
      G.add_edge(node, target)


def update(G):
  # Randomly choose a node
  random_node = random.choice(list(G.nodes))
  opinion_update(G, random_node)
  delete_tie(G, random_node)
  add_tie(G, random_node)
